from keras.callbacks import *
from keras.optimizers import *
import sys,os

sys.path.append('..')
sys.path.append('../..')

"""The image format of all models is "channels_last"""
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# ! <<<
# from Stdlib.paraClass import *
from paraClass import *
# ! >>>
import Unet as Unet
from diceLoss import *
import dataGenerator as DGen
# ! <<< Outdated
import tensorflow as tf
# ! >>>


def ontrain_3D(organ,inputshape,imgpath,labelpath,traintxtpath,testtxtpath,batchsize,modelprefix,sampleparas,normparas,savepath):
    model = Unet.get_3D_Unet(inputshape,sampleparas.nclass,activation='sigmoid',axis=-1)
    # ! <<< Outdated
    # model.compile(optimizer = Adam(lr = 1e-3,clipnorm=1.), loss = dice_coef_loss, metrics = [dice_coef])
    model.compile(optimizer=tf.keras.optimizers.Adam(lr = 1e-3,clipnorm=1.), loss=dice_coef_loss, metrics=[dice_coef])
    # ! >>>
    model_savepath = savepath + "/%s-%s-{epoch:02d}-{dice_coef:.4f}-{val_dice_coef:.4f}.hdf5"%(organ,modelprefix)
    model_checkpoint = ModelCheckpoint(model_savepath, monitor='val_loss',verbose=1, save_best_only=True)
    model_reducelr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, verbose=0, mode='auto', epsilon=0.000001, cooldown=0, min_lr=0)
    print('Fitting model...')

    model.fit_generator(DGen.online_batch_generator_3D(imgpath,labelpath,batchsize,sampleparas,normparas,txtpath=traintxtpath), \
    validation_data=DGen.online_batch_generator_3D(imgpath,labelpath,batchsize,sampleparas,normparas,txtpath=testtxtpath), \
    validation_steps = 100, steps_per_epoch = 1000, epochs=100, verbose=1, shuffle=True, \
    callbacks=[model_checkpoint,model_reducelr])

if __name__ == '__main__':
    organname = 'Rat_Brain'

    # imgpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resample/Raw'
    # labelpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resample/Mask'
    # savepath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/ModelZoo'
    imgpath = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask_3D-Unet/img'
    labelpath = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask_3D-Unet/mask'
    savepath = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask_3D-Unet/out'

    modelprefix = '3D_Unet'
    inputshape = [16,128,128,1]
    batchsize = 2

    normparas = NormParas()
    normparas.method = "minmax"
    # normparas.lmin = 2000
    # normparas.rmax = 18000

    sampleparas = SampleParas()
    sampleparas.patchdims = [16, 128, 128]
    sampleparas.patchlabeldims = [16, 128, 128]
    sampleparas.imgdim = 3
    sampleparas.nclass = 2
    sampleparas.samplenum = 4

    # meanpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_06_HeadNeck_Seg/Data_VOIs_new/mean_2D.npy'
    # sampleparas.meanvalue = np.load(meanpath)

    # traintxtpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resample/RatBrain_train.txt'
    # testtxtpath = '/media/shuaiw/DataDisk/S_Wang_Research/P_Liming_RatBrain_Extraction/TrainData/Total_Mouth_Rat_Resample/RatBrain_val.txt'
    traintxtpath = '/home/yusongli/_project/shidaoai/task/01_seg/2D-Unet/_conf/train_3d.txt'
    testtxtpath = '/home/yusongli/_project/shidaoai/task/01_seg/2D-Unet/_conf/val_3d.txt'

    ontrain_3D(organname,inputshape,imgpath,labelpath,traintxtpath,testtxtpath,batchsize,modelprefix,sampleparas,normparas,savepath)
